Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[nnx] add support for standalone Variables #4606

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[nnx] add support for standalone Variables #4606

wants to merge 1 commit into from
+584 −198

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Mar 7, 2025 •
edited
Loading

What does this PR do?

Adds support for using Variables outside of graph nodes, meaning they can now be passed directly to transforms, including inside pytrees without a parent graph node, and the ability to call split, merge, and update on them directly.

rngs = nnx.Rngs(0)
w = nnx.Param(jax.random.normal(rngs(), (2, 3)))
b = nnx.Param(jnp.zeros((3,)))
count = nnx.Variable(jnp.array(0))

@nnx.jit
def linear(w, b, count, x):
  count += 1
  return x @ w + b[None]

x = jax.random.normal(rngs(), (1, 2))
y = linear(w, b, count, x)

assert count.value == 1
assert y.shape == (1, 3)

@cgarciae cgarciae force-pushed the nnx-standalone-variables branch 3 times, most recently from 21852a6 to fc95f37 Compare March 8, 2025 00:59
@cgarciae cgarciae marked this pull request as ready for review March 8, 2025 01:04
@cgarciae cgarciae force-pushed the nnx-standalone-variables branch 4 times, most recently from 0100223 to ff38c52 Compare March 10, 2025 05:57
@github-advanced-security
Copy link

This pull request sets up GitHub code scanning for this repository. Once the scans have completed and the checks have passed, the analysis results for this pull request branch will appear on this overview. Once you merge this pull request, the 'Security' tab will show more code scanning analysis results (for example, for the default branch). Depending on your configuration and choice of analysis tool, future pull requests will be annotated with code scanning analysis results. For more information about GitHub code scanning, check out the documentation.

@cgarciae cgarciae force-pushed the nnx-standalone-variables branch from ff38c52 to 4ebc477 Compare March 10, 2025 07:31
@cgarciae cgarciae force-pushed the nnx-standalone-variables branch from 4ebc477 to 3f5354d Compare March 10, 2025 08:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@cgarciae